{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import random\n",
    "from typing import List\n",
    "\n",
    "\n",
    "def generate_arithmetic_expression(num: int):\n",
    "    \"\"\"\n",
    "    num: 几个操作符\n",
    "\n",
    "        prompt_template:\n",
    "        你是一名擅长数学运算的助手，负责逐步推理并解决四则运算问题。请按照以下步骤进行：\n",
    "\n",
    "        1. 阅读并理解问题。\n",
    "        2. 分步计算，逐步解决问题。\n",
    "        3. 给出最终的结果。\n",
    "        4. 按照 JSON 格式输出结果，包括：\n",
    "        - reason: 详细的推理过程。\n",
    "        - infer: 最终的计算结果。\n",
    "\n",
    "        问题：{question}\n",
    "        请给出分析和结果。\n",
    "    \"\"\"\n",
    "    # 定义操作符和数字范围，除法\n",
    "    operators = [\"+\", \"-\", \"*\"]\n",
    "    expression = (\n",
    "        f\"{random.randint(1, 100)} {random.choice(operators)} {random.randint(1, 100)}\"\n",
    "    )\n",
    "    num -= 1\n",
    "    for _ in range(num):\n",
    "        expression = f\"{expression} {random.choice(operators)} {random.randint(1, 100)}\"\n",
    "    result = eval(expression)\n",
    "    expression = expression.replace(\"*\", \"x\")\n",
    "    return expression, result\n",
    "\n",
    "\n",
    "def trans2llm_dataset(\n",
    "    texts: List[str],\n",
    "    labels: List[str],\n",
    "    output_file,\n",
    "    instruction=\"\",\n",
    "    prompt_template=\"\",\n",
    "    replace_kw=\"\",\n",
    "):\n",
    "\n",
    "    data = []\n",
    "    for text, label in zip(texts, labels):\n",
    "        if replace_kw and prompt_template:\n",
    "            text = prompt_template.replace(replace_kw, text)\n",
    "\n",
    "        d = {\n",
    "            \"instruction\": instruction,\n",
    "            \"input\": text,\n",
    "            \"output\": label,\n",
    "        }\n",
    "        data.append(d)\n",
    "\n",
    "    with open(output_file, \"w\", encoding=\"utf-8\") as f:\n",
    "        json.dump(data, f, ensure_ascii=False, indent=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_template = \"\"\"\n",
    "    你是一名擅长数学运算的助手，负责逐步推理并解决四则运算问题。请按照以下步骤进行：\n",
    "\n",
    "    1. 阅读并理解问题。\n",
    "    2. 分步计算，逐步解决问题。\n",
    "    3. 给出最终的结果。\n",
    "    4. 按照 JSON 格式输出结果，包括：\n",
    "    - reason: 详细的推理过程。\n",
    "    - infer: 最终的计算结果。\n",
    "\n",
    "    问题：{question}\n",
    "    请给出分析和结果。\n",
    "    \"\"\".strip()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "texts = []\n",
    "labels = []\n",
    "\n",
    "for _ in range(100):\n",
    "    text, label = generate_arithmetic_expression(2)\n",
    "    texts.append(text)\n",
    "    labels.append(str(label))\n",
    "\n",
    "trans2llm_dataset(\n",
    "    texts=texts,\n",
    "    labels=labels,\n",
    "    output_file=\"calculate.json\",\n",
    "    prompt_template=prompt_template,\n",
    "    replace_kw=\"{question}\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llm_beta",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
